#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
拦截管理模块
"""

import logging
import threading
import time
from datetime import datetime
from db.mongo import get_db

logger = logging.getLogger(__name__)

# 全局拦截规则
# 格式: {session_id: {service_name: [url_patterns]}}
intercept_rules = {}

# 拦截状态
# 格式: {session_id: {service_name: {intercepted_count: int, last_intercepted: datetime}}}
intercept_status = {}

# 线程锁
lock = threading.Lock()

def set_intercept_rules(session_id, rules):
    """
    设置拦截规则
    
    Args:
        session_id: 会话ID
        rules: 拦截规则列表，格式为[{service_name: str, url_pattern: str}]
        
    Returns:
        dict: 操作结果
    """
    with lock:
        # 初始化会话的拦截规则
        if session_id not in intercept_rules:
            intercept_rules[session_id] = {}
        
        # 初始化会话的拦截状态
        if session_id not in intercept_status:
            intercept_status[session_id] = {}
        
        # 处理每条规则
        for rule in rules:
            service_name = rule.get('service_name')
            url_pattern = rule.get('url_pattern')
            
            if not service_name or not url_pattern:
                continue
            
            # 添加规则
            if service_name not in intercept_rules[session_id]:
                intercept_rules[session_id][service_name] = []
            
            intercept_rules[session_id][service_name].append(url_pattern)
            
            # 初始化状态
            if service_name not in intercept_status[session_id]:
                intercept_status[session_id][service_name] = {
                    'intercepted_count': 0,
                    'last_intercepted': None
                }
    
    logger.info(f"设置拦截规则: 会话={session_id}, 规则数={len(rules)}")
    return {
        'session_id': session_id,
        'rules_count': len(rules),
        'status': 'active'
    }

def get_intercept_status(session_id):
    """
    获取拦截状态
    
    Args:
        session_id: 会话ID
        
    Returns:
        dict: 拦截状态
    """
    with lock:
        if session_id not in intercept_status:
            return {
                'session_id': session_id,
                'status': 'inactive',
                'services': []
            }
        
        services_status = []
        for service_name, status in intercept_status[session_id].items():
            services_status.append({
                'service_name': service_name,
                'intercepted_count': status['intercepted_count'],
                'last_intercepted': status['last_intercepted'].isoformat() if status['last_intercepted'] else None
            })
        
        return {
            'session_id': session_id,
            'status': 'active',
            'services': services_status
        }

def clear_intercept_rules(session_id):
    """
    清除拦截规则
    
    Args:
        session_id: 会话ID
        
    Returns:
        dict: 操作结果
    """
    with lock:
        if session_id in intercept_rules:
            del intercept_rules[session_id]
        
        if session_id in intercept_status:
            del intercept_status[session_id]
    
    logger.info(f"清除拦截规则: 会话={session_id}")
    return {
        'session_id': session_id,
        'status': 'inactive'
    }

def should_intercept(session_id, service_name, url):
    """
    检查是否应该拦截请求
    
    Args:
        session_id: 会话ID
        service_name: 服务名称
        url: 请求URL
        
    Returns:
        bool: 是否应该拦截
    """
    with lock:
        if session_id not in intercept_rules:
            return False
        
        if service_name not in intercept_rules[session_id]:
            return False
        
        # 检查URL是否匹配任何模式
        for pattern in intercept_rules[session_id][service_name]:
            if pattern in url or pattern == '*':
                return True
        
        return False

def record_interception(session_id, service_name):
    """
    记录拦截事件
    
    Args:
        session_id: 会话ID
        service_name: 服务名称
    """
    with lock:
        if session_id not in intercept_status:
            return
        
        if service_name not in intercept_status[session_id]:
            return
        
        # 更新拦截计数和时间
        intercept_status[session_id][service_name]['intercepted_count'] += 1
        intercept_status[session_id][service_name]['last_intercepted'] = datetime.utcnow()

def get_response_for_request(session_id, service_name, method, url):
    """
    获取请求对应的响应
    
    Args:
        session_id: 会话ID
        service_name: 服务名称
        method: HTTP方法
        url: 请求URL
        
    Returns:
        dict: 响应数据，如果没有找到则返回None
    """
    try:
        db = get_db()
        
        # 查找匹配的请求
        request = db.requests.find_one({
            'session_id': session_id,
            'service_name': service_name,
            'method': method,
            'url': url
        })
        
        if not request:
            logger.warning(f"找不到匹配的请求: {method} {url} -> {service_name}")
            return None
        
        # 查找对应的响应
        response = db.responses.find_one({
            'session_id': session_id,
            'request_id': request['request_id']
        })
        
        if not response:
            logger.warning(f"找不到匹配的响应: {method} {url} -> {service_name}")
            return None
        
        # 记录拦截事件
        record_interception(session_id, service_name)
        
        # 添加拦截标记
        headers = response.get('headers', {})
        headers['X-Intercepted-By'] = 'KubeRelayer'
        headers['X-Intercept-Time'] = datetime.utcnow().isoformat()
        response['headers'] = headers
        
        return response
    
    except Exception as e:
        logger.error(f"获取响应失败: {e}")
        return None
